#coding=utf-8
"""
小块检测
"""
import json
import os
import glob
import sys
import shutil
import cv2
import base64

import numpy as np
import cv2
import requests

class Config:
    def __init__(self):
        self.XIAOKUAI_ASSEMBLE_URL = "http://yj-xiaokuai-ks.haofenshu.com/model_struct_assmble_base64"
        self.load_from_local = False

config = Config()


def is_include(big_rect, small_rect):
    b_x = big_rect[0]
    b_y = big_rect[1]
    b_x2 = b_x + big_rect[2]
    b_y2 = b_y + big_rect[3]
    s_x = small_rect[0]
    s_y = small_rect[1]
    s_x2 = s_x + small_rect[2]
    s_y2 = s_y + small_rect[3]
    if s_x >= b_x and s_y >= b_y  and s_x2 <= b_x2 and s_y2 <= b_y2:
        return  True
    else:
        return False

def encode_base64(img, ext):
    """
    np.ndarray->二进制数据-->base64(str)
    :param img: np.array
    :param ext: extendsion, ".png" or ".jpg"
    :return:
    """
    img_encode = cv2.imencode(ext, img)[1]

    im_base64 = base64.b64encode(img_encode)
    im_base64 = str(im_base64, encoding="utf-8")

    return im_base64


class Rect:
    """矩形结构"""

    def __init__(self, vec4int, pattern=None):
        """
        :param vec4int: int型四元组列表
        :param pattern: 默认为[xmin, ymin, xmax, ymax]
        """
        if pattern == "wh":
            xmin, ymin, w, h = vec4int
            xmax = xmin + w
            ymax = ymin + h
        else:
            xmin, ymin, xmax, ymax = vec4int
        self.xmin = xmin
        self.ymin = ymin
        self.xmax = xmax
        self.ymax = ymax

    def get_area(self):
        """
        计算矩形面积
        """
        w = self.xmax - self.xmin
        h = self.ymax - self.ymin
        area = w * h

        return area

def is_same_row_pure(vec4int1, vec4int2):
    """
    判断两个矩形是否是"同一行",不考虑不同列的情况
    :param vec4int1: int型四元组列表
    :param vec4int2: int型四元组列表
    """
    flag = False

    rect1 = Rect(vec4int1)
    rect2 = Rect(vec4int2)

    rect1_h = rect1.ymax - rect1.ymin
    rect2_h = rect2.ymax - rect2.ymin

    thresh = rect1_h if rect1_h < rect2_h else rect2_h
    thresh /= 2  # 阈值更严格点

    rect1_mid_y = rect1.ymin + rect1_h // 2
    rect2_mid_y = rect2.ymin + rect2_h // 2

    if abs(rect1_mid_y - rect2_mid_y) <= thresh:
        flag = True

    return flag

def fill_white_rects(image, rects):
    """
    填白图像的指定区域
    :param image:
    :param rects: 区域坐标列表,xyxy
    """
    img_h, img_w = image.shape[:2]

    for rect in rects:
        x_min, y_min, x_max, y_max = rect
        x_min -= 5   # 适当放大一点范围
        y_min -= 5
        x_max += 10
        y_max += 10

        if x_max <= img_w and y_max <= img_h:
            image[y_min: y_max, x_min: x_max] = 255

    return image


def get_kaohao_bound_rects(kaohao_obj_data):
    """
    根据服务返回结果提取考号区域坐标
    """
    kaohao_bound_rects = []

    school_number = kaohao_obj_data.get("schoolNumber")
    if not school_number:
        return kaohao_bound_rects
    #bounds = school_number["bounds"]
    bounds = school_number.get("bounds")
    if not bounds:
        return kaohao_bound_rects

    kaohao_x = int(bounds["x"])
    kaohao_y = int(bounds["y"])
    kaohao_w = int(bounds["w"])
    kaohao_h = int(bounds["h"])

    if kaohao_w != 0:
        kaohao_bound_rects.append((kaohao_x, kaohao_y, kaohao_x + kaohao_w, kaohao_y + kaohao_h))

    return kaohao_bound_rects


def get_obj_bound_rects(kaohao_obj_data):
    """
    根据服务返回结果提取客观题区域坐标
    """
    obj_bound_rects = []

    obj_items = kaohao_obj_data["obj_item"]
    if obj_items is not None:
        for obj_item in obj_items:
            region = obj_item["region"]
            obj_x = int(region["x"])
            obj_y = int(region["y"])
            obj_w = int(region["w"])
            obj_h = int(region["h"])
            obj_bound_rects.append((obj_x, obj_y, obj_x+obj_w, obj_y+obj_h))

    return obj_bound_rects


def get_kaohao_object_bound_blocks(kaohao_obj_dict):
    """
    获取考号和客观题的外包围框
    按照x1,y1,x2,y2形式返回
    """
    kaohao_bound_rects = []
    obj_bound_rects = []

    if kaohao_obj_dict["code"] != 0:
        return kaohao_bound_rects, obj_bound_rects

    data = kaohao_obj_dict["data"]

    kaohao_bound_rects = get_kaohao_bound_rects(data)
    obj_bound_rects = get_obj_bound_rects(data)

    return kaohao_bound_rects, obj_bound_rects


def get_zone_coord_list(obj_item_list):
    """
    获取每个小区域的坐标列表
    """
    obj_region_coord_list = []

    for obj_item in obj_item_list:
        region = obj_item.get("region")
        x_min, y_min, w, h = region["x"], region["y"], region["w"], region["h"]
        x_max = x_min + w
        y_max = y_min + h
        coord = [x_min, y_min, x_max, y_max]
        obj_region_coord_list.append(coord)
    
    return obj_region_coord_list 

def get_zone_score_list(obj_item_list):
    """
    获取每个小区域的分数列表
    """
    obj_region_score_list = []

    for obj_item in obj_item_list:
        # 2.记录每个小区域的分数列表
        obj_info_list = obj_item.get("obj")
        if obj_item.get("type") == "same":
            obj_per_score = obj_info_list[0].get("tihao").get("score")  # 每道小题的分数
        else:
            obj_per_score = 0

        obj_region_score = float(obj_per_score)*len(obj_info_list)     # 该区域的总分
        obj_region_score_list.append(obj_region_score)

    return obj_region_score_list


def get_zone_tihao_list(obj_item_list):
    """
    记录每个小区域的题号列表
    """
    obj_dict_list = []

    # 3.记录每个小区域的题号列表
    for obj_item in obj_item_list:
        obj_info_list = obj_item.get("obj")

        obj_dict = {}
        for obj_info in obj_info_list:
            parent_value = str(obj_info.get("tihao").get("pvalue"))   # 父题号
            sub_value = str(obj_info.get("tihao").get("value"))    # 子题号

            obj_dict.setdefault(parent_value, []).append(sub_value)

        obj_dict_list.append(obj_dict)

    return obj_dict_list


def num_same_row_block(in_block, obj_blocks):
    """
    计算in_block与obj_blocks中同一行的数量
    """
    same_row_num = 0

    for block in obj_blocks:
        same_row_flag = is_same_row_pure(in_block, block)
        if same_row_flag:
            same_row_num += 1

    return same_row_num


def filter_obj_blocks_by_shape(obj_blocks, img_h):
    """
    过滤掉图片底部的小块检测的干扰, 并且只保留第一列的客观题区域
    """
    if not obj_blocks:
        return obj_blocks

    ret_blocks = []

    #img_h, _ = image.shape[:2]
    for obj_block in obj_blocks:
        xmin, ymin, xmax, ymax = obj_block
        h = ymax - ymin
        # 如果客观题区域在图片img_h的2/3更下面或img_h的1/3更上面
        # ,且高度小于50, 且同行的客观题区域数量为0,过滤
        same_row_num = num_same_row_block(obj_block, obj_blocks)
        if (ymin>img_h//3*2 or ymin<img_h//3) and h<50 and same_row_num<=1:
            continue

        ret_blocks.append(obj_block)

    return ret_blocks

'''
按照x, y, w, h形式返回
'''
def get_kaohao_obj_blocks(kaohao_obj_dict):
    kaohao_rects = []
    obj_rects = []
    data = kaohao_obj_dict["data"]
    school_number = data["schoolNumber"]
    bounds = school_number["bounds"]
    x = int(bounds["x"])
    y = int(bounds["y"])
    w = int(bounds["w"])
    h = int(bounds["h"])
    kaohao_rects.append((x, y, w, h))

    obj_items = data["obj_item"]
    if obj_items != None:
        for obj_item in obj_items:
            region = obj_item["region"]
            x = int(region["x"])
            y = int(region["y"])
            w = int(region["w"])
            h = int(region["h"])
            obj_rects.append((x, y, w, h))

    return kaohao_rects, obj_rects


def delete_extern_obj_in_col(image, each_obj_blocks, column_blocks):
    for i in range(0,len(column_blocks)):
        block = column_blocks[i]
        rect = [block[0], block[1], block[2] - block[0], block[3] - block[1]]
        obj_blocks = each_obj_blocks[i]
        for obj_rect in obj_blocks:
            if is_include(rect, obj_rect):
                if obj_rect[0] - block[0] > 1000 or block[2] - obj_rect[2] > 1000:
                    # 分列有问题的话，删除左半部分
                    image[obj_rect[1]: obj_rect[3], obj_rect[0]: obj_rect[2]] = 255
                else:
                    image[obj_rect[1]: obj_rect[3], block[0]:block[2]] = 255

        # 客观题之前的区域都可以先抹掉,特别是注意事项一块
        # print("客观题坐标列表obj_blocks: ", obj_blocks)
        obj_blocks = filter_obj_blocks_by_shape(obj_blocks, image)
        if obj_blocks:  # 可能该页没有客观题
            obj_blocks_arr = np.array(obj_blocks)
            obj_xmin = np.min(obj_blocks_arr[:, 0])
            obj_ymin = np.min(obj_blocks_arr[:, 1])
            obj_xmax = np.max(obj_blocks_arr[:, 0]+obj_blocks_arr[:, 2])
            obj_ymax = np.max(obj_blocks_arr[:, 1]+obj_blocks_arr[:, 3])
            #print(obj_ymin, obj_ymax, obj_xmin, obj_xmax)
            #print(image[: obj_ymax, obj_xmin: obj_xmax].shape)
            image[: obj_ymax, : obj_xmax] = 255
    return image


def delete_kaohao_obj_blocks(image, kaohao_obj_dict):
    kaohao_rects, obj_rects = get_kaohao_obj_blocks(kaohao_obj_dict)
    image = fill_white_rects(image, kaohao_rects)
    image = fill_white_rects(image, obj_rects)
    return image


def get_delete_kaohao_obj_blocks(image, kaohao_obj_dict):
    kaohao_bound_rects, obj_bound_rects = get_kaohao_object_bound_blocks(kaohao_obj_dict)
    image = fill_white_rects(image, kaohao_bound_rects)
    image = fill_white_rects(image, obj_bound_rects)
    return kaohao_bound_rects, obj_bound_rects,image


class XiaoKuaiResult:
    """
    小块检测结果
    """
    def __init__(self, kaohao_obj_str, image):
        self.kaohao_obj_str = kaohao_obj_str
        try:    # 如果小块服务有问题,kaohao_obj_str为空字符
            self.kaohao_obj_dict = json.loads(kaohao_obj_str)
        except Exception:
            self.kaohao_obj_dict = {"code": 1}

        self.kaohao_bound_rects = []
        self.obj_bound_rects = []
        self.no_obj_image = image
        self.barcode_bound = [0,0,0,0]
        self.process()

    def process(self):
        """
        分别获取考号区域和客观题区域并填白图片
        """
        self.kaohao_bound_rects, self.obj_bound_rects =\
            get_kaohao_object_bound_blocks(self.kaohao_obj_dict)
        if not config.load_from_local:
            self.no_obj_image = fill_white_rects(self.no_obj_image, self.kaohao_bound_rects)
            self.no_obj_image = fill_white_rects(self.no_obj_image, self.obj_bound_rects)

    def get_kaohao_bounds(self):
        """获取考号"""
        return self.kaohao_bound_rects

    def get_xiaokuai_dict(self):
        """获取小块结果"""
        #print("self.kaohao_obj_dict: ", self.kaohao_obj_dict)
        kaohao_obj_dict = self.kaohao_obj_dict
        code = kaohao_obj_dict.get("code")
        if not code:
            obj_item = kaohao_obj_dict.get("data").get("obj_item")
            schoolNumber = kaohao_obj_dict.get("data").get("schoolNumber")
        else:
            obj_item = []
            schoolNumber = {}

        return obj_item, schoolNumber
    
    def get_obj_list(self):
        '''
        遍历结果，从中找出客观题的题号和选项数，
        return obj_question_list [(tihao, option_num),(tihao, option_num)]
        '''
        obj_items, _ = self.get_xiaokuai_dict()

        obj_question_list = []
        for item in obj_items:
            objs = item["obj"]

            for obj in objs:
                tihao = obj["tihao"]["value"]
                item_pos = obj["item_pos"]
                splitd_item_pos = item_pos.split(";")
                # 要减去1，因为字符串的末尾也有一个;号，会比实际的选项数多1
                option_num = len(splitd_item_pos) - 1 
                obj_question_list.append((tihao, option_num))
        return obj_question_list


    def get_obj_bounds(self):
        """获取客观题"""
        obj_bound_rects = filter_obj_blocks_by_shape(self.obj_bound_rects, self.no_obj_image.shape[0])
        return obj_bound_rects
    def is_on_four_tee(self,h_img, w_img, x,y,w,h, thresh=450):
        flag = False
        if x + w < thresh and y + h < thresh:
            #tl
            flag = True
        elif x > w_img - thresh and y + h < thresh:
            #tr
            flag = True
        elif x + w < thresh and y > h_img - thresh:
            #dl
            flag = True
        elif x > w_img - thresh and y > h_img - thresh:
            #dr
            flag = True
        return flag


    def get_barcode_bound(self):
        """获取barcode"""
        kaohao_obj_dict = self.kaohao_obj_dict
        code = kaohao_obj_dict.get("code")
        self.barcode_bound = [0,0,0,0]
        if len(self.no_obj_image.shape) == 2:
            h_img, w_img = self.no_obj_image.shape
        else:
            h_img, w_img, _ = self.no_obj_image.shape
        if not code:
            barcode = kaohao_obj_dict.get("data").get("barcode_zone")
            if barcode:
                data = barcode.split(';')[0]
                item = data.split(',')
                if len(item)==4:
                    x,y,w,h = int(item[0]),int(item[1]),int(item[2]),int(item[3])
                    flag = self.is_on_four_tee(h_img, w_img, x,y,w,h)
                    if flag:
                        self.barcode_bound = [0, 0, 0, 0]
                    else:
                        #self.barcode_bound = [x,y,x+w,y+h]
                        self.barcode_bound = [x,y,w,h]
        return self.barcode_bound

    def get_no_obj_image(self):
        """获取填白后的图片"""
        return self.no_obj_image

    def get_kaohao_obj_str(self):
        return self.kaohao_obj_str

def get_assmble_xiaokuai_result_base64(img, subject_id, page_idx, is_unit_test=False):
    """
    调用小块服务
    """
    ret_text = ""

    mat_base64 = encode_base64(img, ".png")

    post_dict = dict()
    post_dict['client'] = os.getpid()
    post_dict["mat_base64"] = mat_base64
    post_dict["id"] = subject_id
    post_dict["w"] = img.shape[1]
    post_dict["h"] = img.shape[0]
    post_dict["page_index"] = page_idx
    resquest_data = json.dumps(post_dict)

    try:
        response = requests.post(config.XIAOKUAI_ASSEMBLE_URL, data=resquest_data, timeout=30)
    except Exception as err:
        pass
        # if not is_unit_test:
        #     log_server.logging("request to XIAOKUAI_URL: {} error, {}".format(
        #         config.XIAOKUAI_ASSEMBLE_URL, err))
    else:
        if response.status_code == 200:
            ret_text = response.text
        else:
            pass
            # if not is_unit_test:
            #     log_server.logging("request to XIAOKUAI_URL status_code is {},\
            #         reason is {}".format(response.status_code, response.reason))

    return ret_text


def get_xiaokuai_result(image, img_path, subject_id, page_idx, is_unit_test=False):
    """
    获取小块检测的结果
    考号、客观题、涂抹后的图像
    坐标按照x1,y1,x2,y2保存
    :param image:
    :param subject_id: 考试id
    :param log_server: is_unit_test为False时才有意义
    :param is_unit_test: 是否为单元测试
    :return:
    """
    kaohao_obj_str = get_assmble_xiaokuai_result_base64(image, subject_id, page_idx, is_unit_test)
    ##print("kaohao_obj_str: ", kaohao_obj_str)

    # 本地加载小块结果
    # xiaokuai_result_path = img_path.replace("../struct_std", "debug_result_0311")
    # xiaokuai_result_path = xiaokuai_result_path.replace(".png", "_xiaokuai.txt")
    #print("xiaokuai_result_path ", xiaokuai_result_path)
    # with open(xiaokuai_result_path, encoding="utf-8") as fr:
        # kaohao_obj_str = fr.read()


    xiaokuai_result = XiaoKuaiResult(kaohao_obj_str, image)
    no_obj_image = xiaokuai_result.get_no_obj_image()

    return xiaokuai_result, no_obj_image

def is_ModelAssmbleJson_ok(jsonFile):
    flag =False
    try:  # 如果小块服务有问题,kaohao_obj_str为空字符
        json_result = json.loads(jsonFile)
        struct = json_result['data']
        kaohao_item = struct['schoolNumber']
        obj_cluster = struct['obj_item']
        kaohao = kaohao_item['bounds']
        if obj_cluster != None:
            flag =True
        elif kaohao != None:
            w = kaohao['w']
            if w != 0:
                flag = True
    except Exception:
        flag = False

    return flag

'''
获取考号位置和客观题题号
'''
def Get_key_info_ModelAssmbleJson(jsonFile):
    struct = jsonFile['data']
    # print('strcut:',struct)
    kaohao = []
    obj_v = []

    if struct == "":
        return [],[]
    else:
        kaohao_item = struct['schoolNumber']
        if kaohao_item != None:
            bounds = kaohao_item['bounds']
            if bounds != None:
                x = bounds['x']
                y = bounds['y']
                w = bounds['w']
                h = bounds['h']
                kaohao =[x,y,x+w,y+h]
        obj_cluster = struct['obj_item']
        if obj_cluster != None:
            for cluster in obj_cluster:
                obj_item = cluster['obj']
                if obj_item != None:
                    for item in obj_item:
                        tihao = item['tihao']
                        tihao_value = tihao['value']
                        obj_v.append(tihao_value)
    return kaohao,obj_v

def is_rect_like(box1, box2,ratio=0.7):
    """
    :param box1:[x1,y1,x2,y2] 左上角的坐标与右下角的坐标
    :param box2:[x1,y1,x2,y2]
    :return:
    """
    is_like = False

    width1 = abs(box1[2] - box1[0])
    height1 = abs(box1[1] - box1[3])
    width2 = abs(box2[2] - box2[0])
    height2 = abs(box2[1] - box2[3])
    x_max = max(box1[0], box1[2], box2[0], box2[2])
    y_max = max(box1[1], box1[3], box2[1], box2[3])
    x_min = min(box1[0], box1[2], box2[0], box2[2])
    y_min = min(box1[1], box1[3], box2[1], box2[3])
    iou_width = x_min + width1 + width2 - x_max
    iou_height = y_min + height1 + height2 - y_max
    iou = 0
    if iou_width <= 0 or iou_height <= 0:
        iou = 0
    else:
        iou_area = iou_width * iou_height  # 交集的面积
        box1_area = width1 * height1
        box2_area = width2 * height2
        iou = 1.0* iou_area / (box1_area + box2_area -iou_area )
    #print("iou: {}, ratio: {}, box1: {}, box2: {}".format(iou, ratio, box1, box2))
    if iou > ratio:
        is_like =True
    return is_like

def is_arry_like(list_1, list_2,ratio=0.7):

    is_like = False
    '''
    交集和并集
    '''
    size_1 = len(list_1)
    size_2 = len(list_2)


    if size_1 < 1 or size_2 < 1:
        is_like = False
    elif abs(size_1 - size_2) < 3 and size_1 < 3:
        is_like = True
    else:
        cross = [val for val in list_1 if val in list_2]
        all = (list(set(list_1).union(set(list_2))))
        iou = 1.0 * len(cross) / len(all)
        if iou > ratio:
           is_like = True
    return is_like

def is_xiaokuai_json_near_same(json_1, json_2):
    is_same = False

    kaohao, obj_v = Get_key_info_ModelAssmbleJson(json_1)
    kaohao_2, obj_v_2 = Get_key_info_ModelAssmbleJson(json_2)

    kaohao_like = is_rect_like(kaohao, kaohao_2)
    # print('kaohao_like: ',kaohao_like,kaohao,kaohao_2)
    if kaohao_like:
        is_same = True
    else:
        obj_like = is_arry_like(obj_v, obj_v_2)
        # print('obj_like: ', obj_like, obj_v, obj_v_2)
        if obj_like:
            is_same = True
    return is_same


def get_page_number_xiaokuai(mat_list, is_blank_list,subject_id, is_unit_test=False):
    '''
    利用小块检测结果来判断学生图页码数量
    输入的是base64数组，传出一个学生页码数，和每张图的小块检测结果字符串
    返回0, [] 表示检测失败
    '''

    size = len(mat_list)
    size_b = len(is_blank_list)
    if size < 2:
        return 0, []
    if size != size_b:
        return 0, []
    page_number = 0
    xiaokuai_result = []

    image = mat_list[0]

    page_idx = 0
    result_0 = get_assmble_xiaokuai_result_base64(image, subject_id, page_idx, is_unit_test)
    flag = is_ModelAssmbleJson_ok(result_0)
    if not flag:
        return 0, []
    xiaokuai_result_class = XiaoKuaiResult(result_0, image)
    xiaokuai_result.append(xiaokuai_result_class)
    json_result_0 = json.loads(result_0)

    for i in range(1, size):
         image = mat_list[i]
         is_blank = is_blank_list[i]
         if is_blank:
             continue
         result = get_assmble_xiaokuai_result_base64(image, subject_id, 0, is_unit_test)
         xiaokuai_result_class = XiaoKuaiResult(result, image)
         xiaokuai_result.append(xiaokuai_result_class)
         flag_i = is_ModelAssmbleJson_ok(result)
         if not flag_i:
             continue
         json_result  = json.loads(result)
         is_same = is_xiaokuai_json_near_same(json_result_0,json_result)
         if is_same:
             page_number = i
             break
    return page_number,xiaokuai_result


def test_xiaokuai_bound_result_with_draw(image_path, out_dir):
    """
    小块结果可视化
    """
    if os.path.exists(out_dir):
        shutil.rmtree(out_dir)
    os.makedirs(out_dir)

    green = (0, 255, 0)
    blue = (255, 0, 0)
    purple= (255,0,255)

    image = cv2.imread(image_path, 0)

    base_name = os.path.basename(image_path)
    out_path = os.path.join(out_dir, base_name)

    xiaokuai_result,no_obj_image = get_xiaokuai_result(image, "987", "", True)
    kaohao_bound_rects = xiaokuai_result.get_kaohao_bounds()
    obj_bound_rects = xiaokuai_result.get_obj_bounds()
    barcode_bound  = xiaokuai_result.get_barcode_bound()
    obj_item, schoolNumber = xiaokuai_result.get_xiaokuai_dict()
    #x1,y1,x2,y2 = barcode_bound
    #print('barcode_bound ',x1,y1,x2,y2)
    print('obj_bound_rects : ', obj_bound_rects)
    print('barcode_bound :',barcode_bound)
    # print('obj_item : ', obj_item)
    # print('schoolNumber : ', schoolNumber)
    for item in obj_item:
        print(item)
    corlor = cv2.imread(image_path, 3)
    cv2.rectangle(corlor, (barcode_bound[0], barcode_bound[1]), (barcode_bound[2], barcode_bound[3]), color=purple, thickness=2)
    for rect in kaohao_bound_rects:
        cv2.rectangle(corlor, (rect[0], rect[1]), (rect[2], rect[3]), color=blue, thickness=2)
    for rect in obj_bound_rects:
        cv2.rectangle(corlor, (rect[0], rect[1]), (rect[2], rect[3]), color=green, thickness=3)

    cv2.imwrite(out_path, corlor)


if __name__ == "__main__":
    src_path = "./data/"
    out_dir = "./output/"

    files = glob.glob(src_path+'/*.png')+glob.glob(src_path+'/*.jpg')
    files = files[:1]
    print('img number :', len(files))

    for img_path in files:
        test_xiaokuai_bound_result_with_draw(img_path, out_dir)
